# -*- coding: utf-8 -*-
# -------------------------------------------------
# File Name： tsp_dataset
# Description :
# Author : lirui
# date： 2022/5/31
# -------------------------------------------------
# Change Activity:
# 2022/5/31:
# -------------------------------------------------
import torch
from torch.utils.data import Dataset

from ....common.misc.log_util import logger


class TSPDataset(Dataset):
    """
    TSP training dataset.
    """

    def __init__(self, points_size=50, num_samples=10000):
        """


        Args:
            points_size: the number of points per graph.
            num_samples: total number of training samples.
        """
        self.points_size = points_size
        self.num_samples = num_samples
        self.data = self._build_data()  # shape=(num_samples,points_size,2)

    def _build_data(self):
        """
        build training data, the value of data is greater than 0 and less than 1.

        Returns:
            list: training data.
        """
        logger.debug(f'start creating TSP training data, the points size is {self.points_size}, '
                     f'and number of samples is {self.num_samples}.')
        # data = [torch.FloatTensor(self.points_size, 2).uniform_() for _ in range(self.num_samples)]
        data = torch.rand(size=(self.num_samples, self.points_size, 2))
        logger.debug(f'create TSP training data completed.')
        return data

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx]
